
from __future__ import division

import torch 
import torch.nn as nn
import torch.nn.functional as F 
from torch.autograd import Variable
import numpy as np
import cv2 
# import matplotlib.pyplot as plt
from utils.bbox import bbox_iou
import time

def count_parameters(model):
    return sum(p.numel() for p in model.parameters())

def count_learnable_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def convert2cpu(matrix):
    if matrix.is_cuda:
        return torch.FloatTensor(matrix.size()).copy_(matrix)
    else:
        return matrix

def get_offset(idx,num_anchors,grid_size):

    idx_test = idx[:,1] - idx[:,1]%num_anchors   ### Subtract anchor num offset

    x_offset = (idx_test/num_anchors)%grid_size

    idx_test_1 = idx_test/num_anchors - x_offset

    y_offset = idx_test_1/grid_size

    return x_offset.unsqueeze(1),y_offset.unsqueeze(1)

def advanced_indexing(tensor, index):
    if isinstance(index, tuple):
        adv_loc = []
        for i, el in enumerate(index):
            if isinstance(el, torch.LongTensor):
                adv_loc.append((i, el))
        if len(adv_loc) < 2:
            return tensor[index]
        
        # check that number of elements in each indexing array is the same
        len_array = [i.numel() for _, i in adv_loc]
        #assert len_array.count(len_array[0]) == len(len_array)
        
        idx = [i for i,_ in adv_loc]
        sizes = [tensor.size(i) for i in idx]
        new_size = [tensor.size(i) for i in range(tensor.dim()) if i not in idx]
        new_size_final = [tensor.size(i) for i in range(tensor.dim()) if i not in idx]

        start_idx = idx[0]
        # if there is a space between the indexes
        if idx[-1] - idx[0] + 1 != len(idx):
            permute = idx + [i for i in range(tensor.dim()) if i not in idx]
            tensor = tensor.permute(*permute).contiguous()
            start_idx = 0
        
        lin_idx = _linear_index(sizes, [i for _, i in adv_loc])
        reduc_size = reduce(mul, sizes)
        new_size.insert(start_idx, reduc_size)
        new_size_final[start_idx:start_idx] = list(adv_loc[0][1].size())

        tensor = tensor.view(*new_size)
        tensor = tensor.index_select(start_idx, lin_idx)
        tensor = tensor.view(new_size_final)
        
        return tensor
        
    else:
        return tensor[index]

def filter_boxes_gpu(prediction,conf = 0.2):

    non_zero_ind =  torch.gt(prediction[:,:,4],conf)  ### Comparing only for cars confidence
    idx = torch.nonzero(non_zero_ind)
    prediction = advanced_indexing(prediction, non_zero_ind).unsqueeze(0)

    return non_zero_ind,idx,prediction

def predict_transform(prediction, inp_dim, anchors, num_classes, CUDA = True):

    start = time.time()

    batch_size = prediction.size(0)
    stride =  inp_dim // prediction.size(2)
    grid_size = inp_dim // stride
    ### num_cls = 1 ##
    # num_classes = 1
    ### End ###
    bbox_attrs = 5 + num_classes
    num_anchors = len(anchors)

    anchors = [(a[0]/stride, a[1]/stride) for a in anchors]

    prediction = prediction.view(batch_size, bbox_attrs*num_anchors, grid_size*grid_size)
    prediction = prediction.transpose(1,2).contiguous()
    prediction = prediction.view(batch_size, grid_size*grid_size*num_anchors, bbox_attrs)

    #Sigmoid the  centre_X, centre_Y. and object confidencce
    prediction[:,:,0] = torch.sigmoid(prediction[:,:,0])
    prediction[:,:,1] = torch.sigmoid(prediction[:,:,1])
    prediction[:,:,4] = torch.sigmoid(prediction[:,:,4])

    non_zero_ind_cpu,idx_cpu,prediction_cpu = filter_boxes_gpu(prediction,0.3)
    anchors = torch.FloatTensor(np.asarray(anchors))

    if CUDA:
        anchors = anchors.cuda()
    
    anchors = anchors.repeat(grid_size*grid_size, 1).unsqueeze(0)
    idx_anchors = idx_cpu[:,1]%3
    idx_test_anchor = idx_cpu[:,1].squeeze()
    anchors = anchors[:,idx_test_anchor]

    x_offset, y_offset = get_offset(idx_cpu,num_anchors,grid_size)
    if CUDA:
        x_offset = x_offset.type(torch.cuda.FloatTensor)
        y_offset = y_offset.type(torch.cuda.FloatTensor)

    x_y_offset = torch.cat((x_offset, y_offset), 1).view(-1,2).unsqueeze(0)

    prediction_cpu[:,:,:2] += x_y_offset
    prediction_cpu[:,:,2:4] = torch.exp(prediction_cpu[:,:,2:4])*anchors

    #Softmax the class scores
    prediction_cpu[:,:,5: 5 + num_classes] = torch.sigmoid((prediction_cpu[:,:, 5 : 5 + num_classes]))
    prediction_cpu[:,:,:4] *= stride
    
    return prediction_cpu

def load_classes(namesfile):
    fp = open(namesfile, "r")
    names = fp.read().split("\n")[:-1]
    return names

def get_im_dim(im):
    im = cv2.imread(im)
    w,h = im.shape[1], im.shape[0]
    return w,h

def unique(tensor):
    tensor_np = tensor.cpu().numpy()
    unique_np = np.unique(tensor_np)
    unique_tensor = torch.from_numpy(unique_np)
    
    tensor_res = tensor.new(unique_tensor.shape)
    tensor_res.copy_(unique_tensor)
    return tensor_res

def bbox_iou_test_batch_1(box1, box2):
    """
    Returns the IoU of two bounding boxes 
    
    
    """
    n = int(box1.shape[0])
    
    box1 = box1.unsqueeze(1)
    box2 = box2.unsqueeze(0).repeat(n, 1, 1)
    
    #Get the coordinates of bounding boxes
    b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,:,0], box1[:,:,1], box1[:,:,2], box1[:,:,3]
    b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,:,0], box2[:,:,1], box2[:,:,2], box2[:,:,3]
    
    #get the corrdinates of the intersection rectangle
    inter_rect_x1 =  torch.max(b1_x1, b2_x1)
    inter_rect_y1 =  torch.max(b1_y1, b2_y1)
    inter_rect_x2 =  torch.min(b1_x2, b2_x2)
    inter_rect_y2 =  torch.min(b1_y2, b2_y2)
    
    inter_rect_x1 = inter_rect_x1.type(torch.cuda.FloatTensor)
    inter_rect_y1 = inter_rect_y1.type(torch.cuda.FloatTensor)
    inter_rect_x2 = inter_rect_x2.type(torch.cuda.FloatTensor)
    inter_rect_y2 = inter_rect_y2.type(torch.cuda.FloatTensor)
    
    #Intersection area
    if torch.cuda.is_available():
            inter_area = torch.max(inter_rect_x2 - inter_rect_x1 + 1,torch.zeros(inter_rect_x2.shape).cuda())*torch.max(inter_rect_y2 - inter_rect_y1 + 1, torch.zeros(inter_rect_x2.shape).cuda())
    else:
            inter_area = torch.max(inter_rect_x2 - inter_rect_x1 + 1,torch.zeros(inter_rect_x2.shape))*torch.max(inter_rect_y2 - inter_rect_y1 + 1, torch.zeros(inter_rect_x2.shape))
    
    #Union Area
    b1_area = (b1_x2 - b1_x1 + 1)*(b1_y2 - b1_y1 + 1)
    b2_area = (b2_x2 - b2_x1 + 1)*(b2_y2 - b2_y1 + 1)
        
    b1_area = b1_area.type(torch.cuda.FloatTensor)
    b2_area = b2_area.type(torch.cuda.FloatTensor)
    inter_area = inter_area.type(torch.cuda.FloatTensor)
    
    iou = inter_area / (b1_area + b2_area - inter_area)
    
    #iou_new = inter_area/b2_area
    
    return iou

def write_results_batch(prediction, confidence, num_classes, nms = True, nms_conf = 0.4):
    conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
    prediction = prediction*conf_mask

    # start = time.time()

    # conf_mask = torch.gt(prediction[:,:,4],confidence)
    # print(conf_mask)
    # total_time = time.time() - start

    # prediction = prediction[conf_mask].unsqueeze(0)

    # total_time = time.time() - start
    # print("Total time for Non zero IND : ", total_time)
    
    box_a = prediction.new(prediction.shape)
    box_a[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)
    box_a[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)
    box_a[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2) 
    box_a[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)
    prediction[:,:,:4] = box_a[:,:,:4]
    
    batch_size = prediction.size(0)
    
    output = prediction.new(1, prediction.size(2) + 1)
    write = False

    for ind in range(batch_size):
        #select the image from the batch
        image_pred = prediction[ind]
        max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)

        max_conf = max_conf.float().unsqueeze(1)
        max_conf_score = max_conf_score.float().unsqueeze(1)
        seq = (image_pred[:,:5], max_conf, max_conf_score)
        image_pred = torch.cat(seq, 1)

        #Get rid of the zero entries

        # non_zero_ind =  torch.gt(image_pred[:,4],0)

        mask_lol = image_pred[:,4].gt(0).unsqueeze(1)

        start = time.time()

        image_pred = image_pred.masked_select(mask_lol).view(-1, 7)  

        total_time_1 = time.time() - start
        # print("Total time for Non zero IND : ", total_time_1)

        #Get the various classes detected in the image
        try:
            img_classes = unique(image_pred[:,-1])
        except:
             continue
        #WE will do NMS classwise
        
        for cls in img_classes:
            #get the detections with one particular class
            cls_mask = image_pred*(image_pred[:,-1] == cls).float().unsqueeze(1)
            class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()

            image_pred_class = image_pred[class_mask_ind].view(-1,7)
        
            #sort the detections such that the entry with the maximum objectness
            #confidence is at the top

            conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]
            image_pred_class = image_pred_class[conf_sort_index]
            idx = image_pred_class.size(0)

            start = time.time()
            
            ious = bbox_iou_test_batch_1(image_pred_class, image_pred_class)

            ###Masking ops
        
            ious_masks = (ious < nms_conf)
            
            ### Tril ###
            
            tril = torch.tril(torch.ones(ious.shape[0], ious.shape[0]))
            tril = tril.to(0,dtype=torch.uint8)
            final = torch.max(ious_masks,tril)
            
            nms_idx = unique((final == 0).nonzero()[:,1])
            # print(nms_idx.shape)
            mask_1 = torch.ones(ious.shape[0])
            mask_1 = mask_1.to(0,dtype=torch.uint8)
            
            mask_1[nms_idx] = 0            
            image_pred_class = image_pred_class[mask_1]

            # total_time_1 = time.time() - start
            # print("Total time for IOU : ", total_time_1)
            
            orig_dim = int(ious.shape[0])
                                
            # print(image_pred_class.shape,cls)
            batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
            seq = batch_ind, image_pred_class
            if not write:
                output = torch.cat(seq,1)
                write = True
            else:
                out = torch.cat(seq,1)
                output = torch.cat((output,out))
    
    return output

def write_results(prediction, confidence, num_classes, nms = True, nms_conf = 0.4 , CUDA = True):
    conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
    prediction = prediction*conf_mask
    

    try:
        ind_nz = torch.nonzero(prediction[:,:,4]).transpose(0,1).contiguous()
    except:
        return 0
    
    
    box_a = prediction.new(prediction.shape)
    box_a[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)
    box_a[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)
    box_a[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2) 
    box_a[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)
    prediction[:,:,:4] = box_a[:,:,:4]
    
    batch_size = prediction.size(0)
    
    output = prediction.new(1, prediction.size(2) + 1)
    write = False

    #### Class Filters [0,1,2,3,4] -> boxes [7 12] ####

    # box_class_filter = [0,1,2,3,4,5,8,13]

    #### Class Filters ends #####

    for ind in range(batch_size):
        #select the image from the batch
        image_pred = prediction[ind]
        
        #Get the class having maximum score, and the index of that class
        #Get rid of num_classes softmax scores 
        #Add the class index and the class score of class having maximum score

        max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)
        # max_conf, max_conf_score = torch.max(image_pred[:,box_class_filter], 1)

        max_conf = max_conf.float().unsqueeze(1)
        max_conf_score = max_conf_score.float().unsqueeze(1)
        seq = (image_pred[:,:5], max_conf, max_conf_score)
        image_pred = torch.cat(seq, 1)

        #Get rid of the zero entries
        non_zero_ind =  (torch.nonzero(image_pred[:,4]))
        
        # image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
        image_pred = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
        
        #Get the various classes detected in the image
        try:
            img_classes = unique(image_pred[:,-1])
        except:
             continue
        #WE will do NMS classwise
        for cls in img_classes:
            start = time.time()
            #get the detections with one particular class
            # cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1)
            cls_mask = image_pred*(image_pred[:,-1] == cls).float().unsqueeze(1)

            class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()

            # image_pred_class = image_pred_[class_mask_ind].view(-1,7)
            image_pred_class = image_pred[class_mask_ind].view(-1,7)

             #sort the detections such that the entry with the maximum objectness
             #confidence is at the top
            conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]
            image_pred_class = image_pred_class[conf_sort_index]
            idx = image_pred_class.size(0)
            
            #if nms has to be done
            if nms:
                #For each detection
                for i in range(idx):
                    #Get the IOUs of all boxes that come after the one we are looking at 
                    #in the loop
                    start_1 = time.time()

                    try:
                        ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:] , CUDA)
                    except ValueError:
                        break
        
                    except IndexError:
                        break

                    total_time_1 = time.time() - start_1
                    # print("Total time for Loop 1 : ", total_time_1)

                    # start_1 = time.time()
                    
                    #Zero out all the detections that have IoU > treshhold
                    iou_mask = (ious < nms_conf).float().unsqueeze(1)
                    image_pred_class[i+1:] *= iou_mask       
                    
                    #Remove the non-zero entries
                    non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
                    image_pred_class = image_pred_class[non_zero_ind].view(-1,7)

                    # total_time_1 = time.time() - start_1
                    # print("Total time for cls Loop 2 : ", total_time_1)                   

            #Concatenate the batch_id of the image to the detection
            #this helps us identify which image does the detection correspond to 
            #We use a linear straucture to hold ALL the detections from the batch
            #the batch_dim is flattened
            #batch is identified by extra batch column
            
            
            batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
            seq = batch_ind, image_pred_class
            if not write:
                output = torch.cat(seq,1)
                write = True
            else:
                out = torch.cat(seq,1)
                output = torch.cat((output,out))

            total_time = time.time() - start
            # print("Total time for cls Loop : ", total_time)
        
    return output

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Mar 24 00:12:16 2018

@author: ayooshmac
"""

def predict_transform_half(prediction, inp_dim, anchors, num_classes, CUDA = True):
    batch_size = prediction.size(0)
    stride =  inp_dim // prediction.size(2)

    bbox_attrs = 5 + num_classes
    num_anchors = len(anchors)
    grid_size = inp_dim // stride

    
    prediction = prediction.view(batch_size, bbox_attrs*num_anchors, grid_size*grid_size)
    prediction = prediction.transpose(1,2).contiguous()
    prediction = prediction.view(batch_size, grid_size*grid_size*num_anchors, bbox_attrs)
    
    
    #Sigmoid the  centre_X, centre_Y. and object confidencce
    prediction[:,:,0] = torch.sigmoid(prediction[:,:,0])
    prediction[:,:,1] = torch.sigmoid(prediction[:,:,1])
    prediction[:,:,4] = torch.sigmoid(prediction[:,:,4])

    
    #Add the center offsets
    grid_len = np.arange(grid_size)
    a,b = np.meshgrid(grid_len, grid_len)
    
    x_offset = torch.FloatTensor(a).view(-1,1)
    y_offset = torch.FloatTensor(b).view(-1,1)
    
    if CUDA:
        x_offset = x_offset.cuda().half()
        y_offset = y_offset.cuda().half()
    
    x_y_offset = torch.cat((x_offset, y_offset), 1).repeat(1,num_anchors).view(-1,2).unsqueeze(0)
    
    prediction[:,:,:2] += x_y_offset
      
    #log space transform height and the width
    anchors = torch.HalfTensor(anchors)
    
    if CUDA:
        anchors = anchors.cuda()
    
    anchors = anchors.repeat(grid_size*grid_size, 1).unsqueeze(0)
    prediction[:,:,2:4] = torch.exp(prediction[:,:,2:4])*anchors

    #Softmax the class scores
    prediction[:,:,5: 5 + num_classes] = nn.Softmax(-1)(Variable(prediction[:,:, 5 : 5 + num_classes])).data

    prediction[:,:,:4] *= stride
    
    
    return prediction


def write_results_half(prediction, confidence, num_classes, nms = True, nms_conf = 0.4):
    conf_mask = (prediction[:,:,4] > confidence).half().unsqueeze(2)
    prediction = prediction*conf_mask
    
    try:
        ind_nz = torch.nonzero(prediction[:,:,4]).transpose(0,1).contiguous()
    except:
        return 0    
    
    box_a = prediction.new(prediction.shape)
    box_a[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)
    box_a[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)
    box_a[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2) 
    box_a[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)
    prediction[:,:,:4] = box_a[:,:,:4]
    
    batch_size = prediction.size(0)
    
    output = prediction.new(1, prediction.size(2) + 1)
    write = False
    
    for ind in range(batch_size):
        #select the image from the batch
        image_pred = prediction[ind]

        #Get the class having maximum score, and the index of that class
        #Get rid of num_classes softmax scores 
        #Add the class index and the class score of class having maximum score
        max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)
        max_conf = max_conf.half().unsqueeze(1)
        max_conf_score = max_conf_score.half().unsqueeze(1)
        seq = (image_pred[:,:5], max_conf, max_conf_score)
        image_pred = torch.cat(seq, 1)
        
        #Get rid of the zero entries
        non_zero_ind =  (torch.nonzero(image_pred[:,4]))
        try:
            image_pred_ = image_pred[non_zero_ind.squeeze(),:]
        except:
            continue
        
        #Get the various classes detected in the image
        img_classes = unique(image_pred_[:,-1].long()).half()
        
        
        
                
        #WE will do NMS classwise
        for cls in img_classes:
            #get the detections with one particular class
            cls_mask = image_pred_*(image_pred_[:,-1] == cls).half().unsqueeze(1)
            class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()
            

            image_pred_class = image_pred_[class_mask_ind]

        
             #sort the detections such that the entry with the maximum objectness
             #confidence is at the top
            conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]
            image_pred_class = image_pred_class[conf_sort_index]
            idx = image_pred_class.size(0)
            
            #if nms has to be done
            if nms:
                #For each detection
                for i in range(idx):
                    #Get the IOUs of all boxes that come after the one we are looking at 
                    #in the loop
                    try:
                        ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
                    except ValueError:
                        break
        
                    except IndexError:
                        break
                    
                    #Zero out all the detections that have IoU > treshhold
                    iou_mask = (ious < nms_conf).half().unsqueeze(1)
                    image_pred_class[i+1:] *= iou_mask       
                    
                    #Remove the non-zero entries
                    non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
                    image_pred_class = image_pred_class[non_zero_ind]
                    
                    
            
            #Concatenate the batch_id of the image to the detection
            #this helps us identify which image does the detection correspond to 
            #We use a linear straucture to hold ALL the detections from the batch
            #the batch_dim is flattened
            #batch is identified by extra batch column
            batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
            seq = batch_ind, image_pred_class
            
            if not write:
                output = torch.cat(seq,1)
                write = True
            else:
                out = torch.cat(seq,1)
                output = torch.cat((output,out))
    
    return output
